#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
HAI实例扩缩容测试脚本
测试100个并发请求的扩缩容过程
"""

import asyncio
import aiohttp
import time
import json
import base64
from datetime import datetime
from typing import List, Dict, Any
import statistics

class HAIScalingTest:
    """HAI扩缩容测试器"""
    
    def __init__(self, api_base_url: str = "http://localhost:8080", concurrent_requests: int = 100):
        self.api_base_url = api_base_url
        self.concurrent_requests = concurrent_requests
        self.test_results: List[Dict[str, Any]] = []
        self.system_metrics: List[Dict[str, Any]] = []
        self.start_time = None
        self.end_time = None
        
    def encode_test_image(self) -> str:
        """编码测试图片为base64"""
        # 使用一个小的测试图片数据
        test_image_data = b"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="
        return f"data:image/png;base64,{base64.b64encode(test_image_data).decode('utf-8')}"
    
    async def get_system_status(self, session: aiohttp.ClientSession) -> Dict[str, Any]:
        """获取系统状态"""
        try:
            async with session.get(f"{self.api_base_url}/api/v1/system/status") as response:
                if response.status == 200:
                    data = await response.json()
                    return {
                        'timestamp': time.time(),
                        'active_servers': data['data']['servers']['servers']['available'],
                        'total_servers': data['data']['servers']['servers']['total'],
                        'queue_length': data['data']['queue']['total_tasks_in_queue'],
                        'scaling_running': data['data']['scaling']['running'],
                        'healthy_apps': data['data']['hai']['real_apps']['healthy_apps'],
                        'uptime': data['data']['system']['uptime']
                    }
        except Exception as e:
            print(f"❌ 获取系统状态失败: {e}")
        return {}
    
    async def submit_task(self, session: aiohttp.ClientSession, task_id: int) -> Dict[str, Any]:
        """提交任务到HAI优化系统"""
        start_time = time.time()
        
        try:
            # 准备任务数据
            task_data = {
                'task_id': f"scaling_test_{task_id}_{int(time.time())}",
                'task_data': {
                    'image_data': self.encode_test_image(),
                    'tile_size': '400',
                    'quality_level': 'high'
                }
            }
            
            async with session.post(
                f"{self.api_base_url}/api/v1/tasks/submit",
                json=task_data,
                timeout=aiohttp.ClientTimeout(total=30)
            ) as response:
                end_time = time.time()
                duration = end_time - start_time
                
                if response.status == 200:
                    result = await response.json()
                    return {
                        'task_id': task_id,
                        'success': True,
                        'start_time': start_time,
                        'end_time': end_time,
                        'duration': duration,
                        'response_data': result
                    }
                else:
                    return {
                        'task_id': task_id,
                        'success': False,
                        'error': f'HTTP {response.status}',
                        'start_time': start_time,
                        'end_time': end_time,
                        'duration': duration
                    }
                        
        except asyncio.TimeoutError:
            return {
                'task_id': task_id,
                'success': False,
                'error': 'Request timeout',
                'start_time': start_time,
                'end_time': time.time(),
                'duration': time.time() - start_time
            }
        except Exception as e:
            return {
                'task_id': task_id,
                'success': False,
                'error': str(e),
                'start_time': start_time,
                'end_time': time.time(),
                'duration': time.time() - start_time
            }
    
    async def monitor_system_metrics(self, session: aiohttp.ClientSession, duration: float):
        """监控系统指标"""
        print("📊 开始监控系统指标...")
        start_time = time.time()
        
        while time.time() - start_time < duration:
            metrics = await self.get_system_status(session)
            if metrics:
                self.system_metrics.append(metrics)
                print(f"⏰ {datetime.now().strftime('%H:%M:%S')} - "
                      f"总服务器: {metrics['total_servers']}, "
                      f"活跃服务器: {metrics['active_servers']}, "
                      f"队列: {metrics['queue_length']}")
            await asyncio.sleep(3)  # 每3秒监控一次
    
    async def run_scaling_test(self):
        """运行扩缩容测试"""
        print(f"🚀 开始{self.concurrent_requests}个并发请求的HAI扩缩容测试")
        print("=" * 60)
        
        self.start_time = time.time()
        
        # 创建HTTP会话
        connector = aiohttp.TCPConnector(limit=200, limit_per_host=100)
        timeout = aiohttp.ClientTimeout(total=60)
        
        async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
            # 启动系统监控
            monitor_task = asyncio.create_task(
                self.monitor_system_metrics(session, 180)  # 监控3分钟
            )
            
            # 等待1秒让监控开始
            await asyncio.sleep(1)
            
            # 记录测试开始前的系统状态
            initial_status = await self.get_system_status(session)
            print(f"📊 测试开始前状态: 总服务器={initial_status.get('total_servers', 0)}, "
                  f"活跃服务器={initial_status.get('active_servers', 0)}, "
                  f"队列={initial_status.get('queue_length', 0)}")
            
            # 发送并发请求
            print(f"🔥 发送{self.concurrent_requests}个并发请求...")
            start_requests_time = time.time()
            
            tasks = [
                self.submit_task(session, i) 
                for i in range(self.concurrent_requests)
            ]
            
            # 等待所有请求完成
            results = await asyncio.gather(*tasks, return_exceptions=True)
            
            end_requests_time = time.time()
            requests_duration = end_requests_time - start_requests_time
            
            print(f"✅ 所有请求完成，耗时: {requests_duration:.2f}秒")
            
            # 继续监控一段时间观察缩容
            print("⏳ 继续监控30秒观察缩容过程...")
            await asyncio.sleep(30)
            
            # 停止监控
            monitor_task.cancel()
            
            # 记录测试结束后的系统状态
            final_status = await self.get_system_status(session)
            print(f"📊 测试结束后状态: 总服务器={final_status.get('total_servers', 0)}, "
                  f"活跃服务器={final_status.get('active_servers', 0)}, "
                  f"队列={final_status.get('queue_length', 0)}")
            
            # 处理结果
            self.test_results = []
            for i, result in enumerate(results):
                if isinstance(result, Exception):
                    self.test_results.append({
                        'task_id': i,
                        'success': False,
                        'error': str(result),
                        'start_time': start_requests_time,
                        'end_time': time.time(),
                        'duration': time.time() - start_requests_time
                    })
                else:
                    self.test_results.append(result)
        
        self.end_time = time.time()
        
        # 分析结果
        self.analyze_results()
    
    def analyze_results(self):
        """分析测试结果"""
        print("\n📊 测试结果分析")
        print("=" * 60)
        
        total_requests = len(self.test_results)
        successful_requests = len([r for r in self.test_results if r['success']])
        failed_requests = total_requests - successful_requests
        
        print(f"总请求数: {total_requests}")
        print(f"成功请求: {successful_requests}")
        print(f"失败请求: {failed_requests}")
        print(f"成功率: {successful_requests/total_requests*100:.1f}%")
        
        if successful_requests > 0:
            durations = [r['duration'] for r in self.test_results if r['success']]
            print(f"\n⏱️ 响应时间统计:")
            print(f"  平均响应时间: {statistics.mean(durations):.3f}秒")
            print(f"  最短响应时间: {min(durations):.3f}秒")
            print(f"  最长响应时间: {max(durations):.3f}秒")
            print(f"  中位数响应时间: {statistics.median(durations):.3f}秒")
        
        # 分析系统扩缩容
        if self.system_metrics:
            print(f"\n🔄 系统扩缩容分析:")
            initial_servers = self.system_metrics[0]['total_servers']
            max_servers = max(m['total_servers'] for m in self.system_metrics)
            final_servers = self.system_metrics[-1]['total_servers']
            
            print(f"  初始服务器数: {initial_servers}")
            print(f"  最大服务器数: {max_servers}")
            print(f"  最终服务器数: {final_servers}")
            print(f"  扩容数量: {max_servers - initial_servers}")
            
            # 找到扩容时间点
            for i, metrics in enumerate(self.system_metrics):
                if metrics['total_servers'] > initial_servers:
                    print(f"  扩容开始时间: {datetime.fromtimestamp(metrics['timestamp']).strftime('%H:%M:%S')}")
                    break
            
            # 找到缩容时间点
            for i, metrics in enumerate(self.system_metrics):
                if metrics['total_servers'] < max_servers:
                    print(f"  缩容开始时间: {datetime.fromtimestamp(metrics['timestamp']).strftime('%H:%M:%S')}")
                    break
    
    def save_results(self, filename: str = None):
        """保存测试结果"""
        if filename is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"hai_scaling_test_results_{timestamp}.json"
        
        results = {
            'test_info': {
                'timestamp': datetime.now().isoformat(),
                'concurrent_requests': self.concurrent_requests,
                'api_base_url': self.api_base_url,
                'start_time': self.start_time,
                'end_time': self.end_time,
                'total_duration': self.end_time - self.start_time if self.end_time and self.start_time else 0
            },
            'test_results': self.test_results,
            'system_metrics': self.system_metrics,
            'analysis': {
                'total_requests': len(self.test_results),
                'successful_requests': len([r for r in self.test_results if r['success']]),
                'failed_requests': len([r for r in self.test_results if not r['success']]),
                'success_rate': len([r for r in self.test_results if r['success']]) / len(self.test_results) * 100 if self.test_results else 0
            }
        }
        
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(results, f, indent=2, ensure_ascii=False)
        
        print(f"\n💾 测试结果已保存到: {filename}")
        return filename

async def main():
    """主函数"""
    print("🎯 HAI扩缩容测试开始")
    print("=" * 60)
    
    # 创建测试器
    tester = HAIScalingTest(concurrent_requests=100)
    
    # 运行测试
    await tester.run_scaling_test()
    
    # 保存结果
    filename = tester.save_results()
    
    print(f"\n🎉 测试完成！结果已保存到: {filename}")

if __name__ == "__main__":
    asyncio.run(main())
